import cv2
import numpy as np

from src import enum
from py_base import tool

def drawCutLine(mask, cutBox, shotBox):
    h, w = mask.shape[:2]
    halfw = int(w/2)
    # 绘制裁切分割线
    cv2.line(mask, (cutBox[0], 0), (cutBox[0], h), (0, 255, 255), 1, lineType=cv2.LINE_AA)
    cv2.line(mask, (cutBox[1], 0), (cutBox[1], h), (0, 255, 255), 1, lineType=cv2.LINE_AA)
    cv2.line(mask, (0, cutBox[2]), (w, cutBox[2]), (0, 255, 255), 1, lineType=cv2.LINE_AA)
    cv2.line(mask, (0, cutBox[3]), (w, cutBox[3]), (0, 255, 255), 1, lineType=cv2.LINE_AA)

    cv2.line(mask, (shotBox[0], 0), (shotBox[0], h), (255, 0, 255), 1, lineType=cv2.LINE_AA)
    cv2.line(mask, (shotBox[1], 0), (shotBox[1], h), (255, 0, 255), 1, lineType=cv2.LINE_AA)

    cv2.line(mask, (halfw, 0), (halfw, h), (0, 0, 255), 1, lineType=cv2.LINE_AA)


def predict_test(img):
    cutBox = list(map(int, enum.setting['cutBox'].split(",")))
    # cropped_image = image[y_start:y_end, x_start:x_end]
    cutBoxImg = img[cutBox[2]:cutBox[3], cutBox[0]:cutBox[1]]

    cutBoxImgGray = cv2.cvtColor(cutBoxImg, cv2.COLOR_BGR2GRAY)

    cutBoxImgHsv = cv2.cvtColor(cutBoxImg, cv2.COLOR_BGR2HSV)
    svSet = tool.getEnumIntArr('sv_offset')
    vSet = svSet[1]
    conf = tool.getEnumIntArrs('dynamicV')
    if conf is not None:
        brightness = int(np.mean(cutBoxImgHsv))
        vConf, imgConf = conf
        if brightness <= imgConf[0]:
            vSet = vConf[0]
        elif brightness >= imgConf[1]:
            vSet = vConf[1]
        else:
            bi = (brightness - imgConf[0]) / (imgConf[1] - imgConf[0])
            vSet = vConf[1] - int((vConf[1] - vConf[0]) * bi)
        tool.putText(img, f"dynamicV:{brightness}, {vSet}", 2)



    cutBoxImgHsv = tool.adjust_hsv(cutBoxImgHsv, svSet[0]/10, vSet/10)

    hsv = np.array(tool.getEnumIntArr('hsv'))
    offset = np.array(tool.getEnumIntArr('offset'))
    hsvMask = cv2.inRange(cutBoxImgHsv, hsv - offset, hsv + offset)

    cutBoxImgGray = cv2.bitwise_and(cutBoxImgGray, cutBoxImgGray, mask=hsvMask)

    shotBox = tool.getEnumIntArr('shotBox')
    shotBoxImg = cutBoxImgGray[:, shotBox[0] - cutBox[0]: shotBox[1] - cutBox[0]]

    cutBoxImgGray = cv2.GaussianBlur(cutBoxImgGray, (5, 5), 8)  # 高斯滤波，cutBoxImg会脱离原图
    # ret, cutBoxImg = cv2.threshold(cutBoxImg, 140, 255, cv2.THRESH_BINARY)

    cannySet = tool.getEnumIntArr('canny_offset')
    cannyMask = cv2.Canny(cutBoxImgGray, cannySet[0], cannySet[1])  # canny 边缘检测,得到二值化图片

    cutBoxImgHsv = cv2.cvtColor(cutBoxImgHsv, cv2.COLOR_HSV2BGR)
    cv2.imshow('cutBoxImgHsv', cutBoxImgHsv)

    # contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # 提取轮廓
    # cntImg = np.zeros(cutBoxImg.shape, dtype=np.uint8)
    # cv2.drawContours(cntImg, contours, -1, (255, 255, 255), 1) # 将轮廓 画到一张纯黑的图上


    # 检测直线
    hfLine = tool.getEnumIntArr('hfLine')
    # 1. ** 阈值（threshold） ** ：霍夫变换算法中累加器中的投票数阈值。只有当累加器中某个（ρ, θ）组合的投票数超过此阈值时，才会被认为是一条直线。
    # 2. ** 最小线段长度（minLineLength） ** ：可以被认为是一条直线的最小长度。较大的值会忽略较短的线段。
    # 3. ** 最大间隔（maxLineGap） ** ：允许将同一条直线上的点连接起来的最大间隔。较大的值会把间隔较远的点连接成一条线段。
    lines = cv2.HoughLinesP(cannyMask, hfLine[0], np.pi / hfLine[1], hfLine[2],
                            minLineLength=hfLine[3], maxLineGap=hfLine[4])

    resLines = []
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if x2 < cannyMask.shape[1] / 2:
                x1 += cutBox[0]
                x2 += cutBox[0]
                y1 += cutBox[2]
                y2 += cutBox[2]
                resLines.append([x1, y1, x2, y2])
                cv2.line(img, (x1, y1), (x2, y2), (255, 0, 0), 2, lineType=cv2.LINE_AA)
                if len(resLines) == 1:
                    flag = -1 if y2 - y1 > 0 else 1
                    cv2.line(img, (x2, y2), (x2, y2 + flag * enum.setting['minGap']), (0, 0, 255), 2, lineType=cv2.LINE_AA)




    # 识别圆形
    circles = []
    shotBoxImg = cv2.GaussianBlur(shotBoxImg, (5, 5), 8)  # 高斯滤波，shotBoxImg会脱离原图
    mask = cv2.Canny(shotBoxImg, cannySet[0], cannySet[1])  # canny 边缘检测,得到二值化图片

    cv2.imshow('cutBoxImgGray', np.hstack((cutBoxImgGray, hsvMask, cannyMask, shotBoxImg, mask)))

    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    shotItemLeft = []
    shotItemRight = []
    h, w = img.shape[:2]
    halfw = int(w/2)
    for contour in contours:
        if len(contour) > 5:  # 确保轮廓点数足够
            x, y, w, h = cv2.boundingRect(contour)  # 外接矩形
            x += shotBox[0]
            y += cutBox[2]
            cv2.rectangle(img, (x, y), (x + w, y + h), (255, 255, 255), 2)
            if 5 < w < 15 and 5 < h < 15:

                if x < halfw and x+w < halfw and halfw - (x+w) < 20: # 框位于左侧，且距离中线较近
                    shotItemLeft.append([x, y, x+w, y+h])
                else:
                    b = x + 5
                    if b > halfw and b + w > halfw and b - halfw < 20:
                        shotItemRight.append([x, y, x+w, y+h])

    shotItemLeft = tool.filter_boxes_by_iou(shotItemLeft, 0.5)
    shotItemRight = tool.filter_boxes_by_iou(shotItemRight, 0.5)

    shotItemRes = [[], []]
    yGap = 10
    for left in shotItemLeft:
        for right in shotItemRight:
            # 左侧框和右侧框y轴对称
            l_x1, l_y1, l_x2, l_y2 = left
            r_x1, r_y1, r_x2, r_y2 = right
            if l_y1 - yGap < r_y1 < l_y1 + yGap and l_y2 - yGap < r_y2 < l_y2 + yGap:
                shotItemRes[0].append(left)
                shotItemRes[1].append(right)



    for i in shotItemRes:
        for item in i:
            x, y, x2, y2 = item
            cv2.rectangle(img, (x, y), (x2, y2), (0, 255, 0), 2)

    coffset = 2
    if len(shotItemRes[0]) >= 2:
        l_x1, l_y1, l_x2, l_y2 = shotItemRes[0][0]
        center = [l_x2 + coffset, l_y2 + coffset]
        circles.append(center)
        cv2.circle(img, center, 2, (0, 0, 255), 2)

    drawCutLine(img, cutBox, shotBox)
    return resLines, circles, img


    # circles = cv2.HoughCircles(shotBoxImg, cv2.HOUGH_GRADIENT, 30, 10, param2=5, maxRadius =80)
    # if circles is not None:
    #     # 将检测的圆画出来
    #     for i in circles[0, :]:
    #         cv2.circle(drawing, (int(i[0]) + cutBox[0], int(i[1]) + cutBox[2]), int(i[2]), (0, 255, 0), 2)  # 画出外圆
    #         cv2.circle(drawing, (int(i[0]) + cutBox[0], int(i[1]) + cutBox[2]), 2, (0, 0, 255), 3)  # 画出圆心



